# load packages
library(tidyverse)
library(brms)
library(tidybayes)
library(hrbrthemes)
library(kableExtra)

fit202 <- read_rds("output/fits/model-01-structure-2-0-2.rds") 

summary(fit202)

df <- read_rds("output/rescaled-data.rds") %>%
  glimpse()

# overall (no res) ----
## create fake data set
fake_df <- crossing(treat_indicator = c(0, 1), awareness = sort(unique(df$awareness))) %>%
  left_join(distinct(select(df, awareness, awareness_std))) %>%
  mutate(treat_indicator_std = treat_indicator - 0.5) %>%
  glimpse()
obs_awareness <- df %>%
  select(awareness)
## simulate quantities of interest
qi_df <- fake_df %>%
  add_epred_draws(fit202, re_formula = NA) %>%
  ungroup() %>% 
  select(-ends_with("_std")) %>%
  group_by(awareness, .draw) %>%
  select(-.chain, -.iteration, -.row) %>%
  spread(treat_indicator, .epred) %>%
  mutate(treatment_effect = `1` - `0`) %>%
  select(-`0`, -`1`) %>% 
  select(awareness, draw = .draw, treatment_effect) %>%
  ungroup() %>%
  glimpse()
# summaries qis
qi_sum_df <- qi_df %>%
  group_by(awareness) %>%
  summarize(post_avg = mean(treatment_effect), 
            post95 = quantile(treatment_effect, 0.95), 
            post05 = quantile(treatment_effect, 0.05)) %>%
  ungroup() %>%
  glimpse()
# plot treatment effect estimates and cis
ggplot(qi_sum_df, aes(x = awareness, y = post_avg, 
                            ymin = post05, ymax = post95)) + 
  geom_ribbon(fill = "grey90") +
  geom_line() + 
  #facet_grid(cols = vars(category), scales = "free_y", space = "free") + 
  scale_x_continuous(labels = scales::percent_format(accuracy = 1)) + 
  labs(title = "Estimates of Treatment Effect for Typical Policy as Awareness Varies",
       subtitle = "Posterior Average and 90% Percentile Interval",
       y = "Treatment Effect", 
       x = "Percent Aware of the Parties' Position") + 
  geom_rug(data = obs_awareness, aes(x = awareness, y = NULL, ymin = NULL, ymax = NULL), sides = "b") + 
  theme_bw() + theme(legend.position = "bottom", legend.title.align = 1)
ggsave("figs/fig3-model-overall.tiff",
       width = 4, height = 3, scale = 1.7)

xlo_tmp <- qi_sum_df$post_avg[qi_sum_df$awareness == max(qi_sum_df$awareness)]
xhi_tmp <- qi_sum_df$post_avg[qi_sum_df$awareness == min(qi_sum_df$awareness)]
x_tmp <- paste0("...and a minimum-to-maximum increase by about ", 
                round(xhi_tmp - xlo_tmp, 2), 
                " units (from a treatment effect of ", 
                round(xlo_tmp, 2),
                " to ", 
                round(xhi_tmp, 2),
                "..."); x_tmp
cat(x_tmp,
    file = "output/qis.txt", append = TRUE)


# by category ----
## create fake data set
fake_df <- crossing(treat_indicator = c(0, 1), 
                    category = sort(unique(df$category))) %>%
  left_join(distinct(select(df, category, awareness, awareness_std))) %>%
  mutate(treat_indicator_std = treat_indicator - 0.5) %>%
  glimpse()
obs_awareness <- df %>%
  select(category, awareness) %>%
  mutate(category = fct_recode(category, 
                               `Economic Policy` = "Economic",
                               `Social Policy` = "Social"))
## simulate quantities of interest
qi_df <- fake_df %>% 
  add_epred_draws(fit202, re_formula = ~ (1 + treat_indicator_std | category)) %>% 
  ungroup() %>% 
  select(-ends_with("_std")) %>%
  group_by(category, awareness, .draw) %>%
  select(-.chain, -.iteration, -.row) %>%
  spread(treat_indicator, .epred) %>%
  mutate(treatment_effect = `1` - `0`) %>%
  select(-`0`, -`1`) %>% 
  select(category, awareness, draw = .draw, treatment_effect) %>%
  ungroup() %>%
    mutate(category = fct_recode(category, 
                                 `Economic Policy` = "Economic",
                                 `Social Policy` = "Social")) %>%
  glimpse()
# summaries qis
qi_sum_df <- qi_df %>%
  group_by(category, awareness) %>%
  summarize(post_avg = mean(treatment_effect), 
            post95 = quantile(treatment_effect, 0.95), 
            post05 = quantile(treatment_effect, 0.05)) %>%
  ungroup() %>%
  group_by(category) %>%
  mutate(min_awareness = min(awareness),
         max_awareness = max(awareness)) %>% 
  filter(awareness <= max_awareness & awareness >=  min_awareness) %>%
  glimpse()

qi_sum_df_unlabeled <- qi_sum_df %>%
  rename(category_star = category) %>%
  glimpse
# plot treatment effect estimates and cis
ggplot(qi_sum_df, aes(x = awareness, y = post_avg, 
                      ymin = post05, ymax = post95)) + 
  geom_ribbon(aes(fill = category), alpha = 0.3) +
  #geom_line(data = qi_sum_df_unlabeled, aes(group = category_star), size = 0.5, color = "grey50") + 
  geom_line(aes(color = category), size = 1.2) + 
  #facet_grid(cols = vars(category), scales = "free_y", space = "free") + 
  scale_x_continuous(labels = scales::percent_format(accuracy = 1)) + 
  labs(title = "Estimates of Treatment Effect for a Typical Policy in Each Category as Awareness Varies",
       subtitle = "Posterior Average and 90% Percentile Interval",
       y = "Treatment Effect", 
       x = "Percent Aware of the Parties' Position") + 
  geom_rug(data = obs_awareness, aes(x = awareness, y = NULL, ymin = NULL, ymax = NULL, color = category), sides = "b") + 
  theme_bw() + theme(legend.position = "bottom", legend.title.align = 1) + 
  geom_line(data = qi_sum_df_unlabeled, aes(color = category_star))

ggsave("figs/fig4-model-by-category.tiff",
       width = 8, height = 5, scale = 1.3)


cats <- gather_draws(fit202, r_category[category, parameter]) %>%
  filter(parameter == "treat_indicator_std") %>%
  pivot_wider(names_from = "category", values_from = ".value") %>%
  mutate(fs_diff = Foreign.Policy - Social,
         fe_diff = Foreign.Policy - Economic,
         es_diff = Economic - Social) %>%
  summarize(fs_p = mean(fs_diff > 0),
            fe_p = mean(fe_diff < 0),
            es_p = mean(es_diff > 0),
            e_fs = mean(fs_diff),
            e_fe = mean(fe_diff),
            e_es = mean(es_diff)) %>%
  ungroup() %>%
  select(contains("_p"), contains("e_")) %>%
  glimpse()

cat("\n\nThe table below shows the category comparisons.
     f :: foreign policy
     s :: social policy
     e :: economic policy
     _p :: p-value for 1 minus 2
     e_ :: expected value of 1 minus 2\n\n", 
    file = "output/qis.txt", append = TRUE)
k <- kable(cats, format = "simple", digits = 2)
cat(paste0(k, collapse =  "\n"),  
    file = "output/qis.txt", append = TRUE)
writeLines(k, con = "output/qis.txt", append = TRUE)
save_kable(k, file = "output/qis.txt", format = "simple")



